
import xml.dom.minidom as Dom
import os
'''此文件不能单独使用，需其它脚本引用，入口为detection2xml（inference_relabel.py中有示例）
filename：图片文件路径
size：图片的宽，长，维度（jpg为3）
boxlist：bboox列表，5个元素：label,xmin,ymin,xmax,ymax'''
class XMLGenerator:
    def __init__(self, xml_name):
        self.xml_name = xml_name
        if not os.path.exists(xml_name):
            self.doc = Dom.Document()
        else:
            with open(xml_name,'r',encoding='utf-8') as f:
                self.doc = Dom.parse(f)
    def createNode(self, node_name):
        return self.doc.createElement(node_name)

    def addNode(self, node, prev_node = None):
        cur_node = node
        if prev_node is not None:
            prev_node.appendChild(cur_node)
        else:
            self.doc.appendChild(cur_node)
        return cur_node

    def setNodeAttr(self, node, att_name, value):
        cur_node = node
        cur_node.setAttribute(att_name, value)

    def setNodeValue(self, cur_node, value):
        node_data = self.doc.createTextNode(value)
        cur_node.appendChild(node_data)

    def genXml(self):
        f = open(self.xml_name, "w")
        f.write(self.doc.toprettyxml(indent = "\t",newl = "\n",encoding='utf-8').decode('utf-8'))
        f.close()


def detection2xml(filename,size,boxlist):
    filejpg = filename.split('/')[-1]
    filejpg = filejpg.split('.')[0] + '.jpg'
    myXMLGenerator = XMLGenerator(filename)
    #xml root node
    if not os.path.exists(filename):
        node_annotation = myXMLGenerator.createNode("annotation")
        myXMLGenerator.setNodeAttr(node_annotation, "verified", "yes")
        myXMLGenerator.addNode(node = node_annotation)

        nodelist = ['folder', 'filename', 'path', 'segmented']
        valuelist = ['detection',filejpg, filejpg, '0']
        for i in range(4):
          node_tmp = myXMLGenerator.createNode(nodelist[i])
          myXMLGenerator.setNodeValue(node_tmp, valuelist[i])
          myXMLGenerator.addNode(node_tmp, node_annotation)

        node_source = myXMLGenerator.createNode("source")
        node_database = myXMLGenerator.createNode("database")
        myXMLGenerator.setNodeValue(node_database, "DDJ")
        myXMLGenerator.addNode(node_database, node_source)
        myXMLGenerator.addNode(node_source, node_annotation)

        node_size = myXMLGenerator.createNode("size")
        nodelist = ['width', 'height', 'depth']
        for i in range(3):
          node_tmp = myXMLGenerator.createNode(nodelist[i])
          myXMLGenerator.setNodeValue(node_tmp, str(size[i]))
          myXMLGenerator.addNode(node_tmp, node_size)
        myXMLGenerator.addNode(node_size, node_annotation)
    else:
        node_annotation = myXMLGenerator.doc.documentElement


    for obj in boxlist:
      node_obj = myXMLGenerator.createNode("object")
      nodelist = ['name', 'pose', 'truncated', 'difficult']
      valuelist = [obj[0],'Unspecified','0', '0']
      for i in range(4):
        node_tmp = myXMLGenerator.createNode(nodelist[i])
        myXMLGenerator.setNodeValue(node_tmp, valuelist[i])
        myXMLGenerator.addNode(node_tmp, node_obj)
      node_box = myXMLGenerator.createNode('bndbox')
      nodelist = ['xmin', 'ymin', 'xmax', 'ymax']
      # if obj[3] > size[0] and obj[4] > size[1]:
      #   continue
      valuelist = [str(obj[1]),str(obj[2]),str(obj[3]),str(obj[4])]
      for i in range(4):
        node_tmp = myXMLGenerator.createNode(nodelist[i])
        myXMLGenerator.setNodeValue(node_tmp, valuelist[i])
        myXMLGenerator.addNode(node_tmp, node_box)
      myXMLGenerator.addNode(node_box, node_obj)
      myXMLGenerator.addNode(node_obj, node_annotation)
    #gen
    myXMLGenerator.genXml()
    with open(filename,'r') as f:
        lines= f.readlines()
        lines=[l for l in lines if l[-2:]!='\t\n']
    with open(filename,'w') as f:
        for l in lines:
            f.write(l)
#usage:detection2xml('filepath',[int(width),int(height),int(3)],[[class_name,xmin,ymin,xmax,ymax],])